import pandas as pd
import numpy as np
from scipy.stats import chi2_contingency
from collections import Counter
from scipy.stats import entropy

def cramers_v(x, y):
    r"""
    计算Cramér's V统计量。
    
    在统计学中，克拉梅尔V系数（常记作克拉梅尔 :math:`\varphi` 或 
    :math:`\varphi_c`）是两个名义变量关联度的度量，给定一个介于 0
    到 1 之间的值。这一方法基于皮尔逊卡方统计改进而来，由哈拉尔德·克拉
    梅尔于 1946 年首次提出。

    其计算方法为：对于变量 :math:`A` 和 :math:`B`，假设 :math:`A`
    有 :math:`k` 个类别，:math:`B` 有 :math:`r` 个类别，我们有一个
    包含频数的 :math:`k \times r` 列联表。 :math:`n_{i,j}` 表示第
    :math:`i` 行第 :math:`j` 列的频数，即 :math:`(A_i, B_j)` 的
    联合频数。

    卡方统计量 :math:`\chi^2` 的计算方法为：

    .. math::

        \chi^2 = \sum_{i,j} 
        \dfrac{\left(n_{ij} - 
        \dfrac{n_{i.} n_{.j}}{n}\right)^2}
        {\dfrac{n_{i.} n_{.j}}{n}}

    其中， 
    
    - :math:`n_{i.}` 是第 :math:`i` 行的总频数
    - :math:`n_{.j}` 是第 :math:`j` 列的总频数
    - :math:`n` 是总样本量。

    克拉梅尔V系数是通过将卡方统计量的平方根除以样本量和最小维数减去 1
    计算所得的：
    
    .. math::

        \varphi_c = \sqrt{\dfrac{\chi^2 / n}{\min(k-1, r-1)}}

    显著性 :math:`p` 值的计算方法与皮尔逊相关性一致，通过 
    :math:`\chi^2` 值和相应的自由度 :math:`(k-1)(r-1)` 查找卡方
    分布表可以确定 :math:`p` 值。克拉梅尔 V 系数提供了一种衡量名义
    变量之间关系强度的方法，值越接近 1 表示关联越强。

    输入参数
    ----------
    x : array-like
        类别型变量。
    y : array-like
        类别型变量。

    返回值
    ----------
    float
        Cramér's V值。
    """
    confusion_matrix = pd.crosstab(x, y)
    chi2 = chi2_contingency(confusion_matrix)[0]
    n = confusion_matrix.sum().sum()
    phi2 = chi2 / n
    r, k = confusion_matrix.shape
    phi2corr = max(0, phi2 - ((k - 1) * (r - 1)) / (n - 1))
    rcorr = r - ((r - 1) ** 2) / (n - 1)
    kcorr = k - ((k - 1) ** 2) / (n - 1)
    return np.sqrt(phi2corr / min((kcorr - 1), (rcorr - 1)))

def theils_u(x, y):
    r"""
    计算Theil's U统计量。
    
    Theil's U 计算公式如下：

    .. math::

    U(x|y) = \frac{S(x) - S(x|y)}{S(x)}

    其中 :math:`S(x)` 为边际熵 (marginal entropy)：

    .. math::

      S(x) = -\sum_{i=1}^{n} P(x_i) \\log P(x_i)

    :math:`S(x|y)` 为条件熵 (conditional entropy)：

    .. math::

      S(x|y) = -\sum_{j=1}^{m} P(y_j) \sum_{i=1}^{n} 
          P(x_i|y_j) \log P(x_i|y_j)
      
    :math:`P(x_i)` 为 x 的概率分布， :math:`P(x_i|y_j)` 为在 y 
    给定的条件下 x 的条件概率分布。

    注意
    ----------
    Theil's U 是不对称的，即 :math:`U(x|y)` 不等于 :math:`U(y|x)`。
    因此，计算时需注意变量顺序。

    示例
    ----------
    
    >>> x = ['a', 'b', 'a', 'a', 'b', 'c']
    >>> y = ['1', '1', '2', '2', '3', '3']
    >>> theils_u(x, y)
    0.6666666666666666

    输入参数
    ----------
    x : array-like
        类别型变量。
    y : array-like
        类别型变量。

    返回值
    ----------
    float
        Theil's U值。
    """
    s_xy = conditional_entropy(x, y)
    x_counter = Counter(x)
    total_occurrences = sum(x_counter.values())
    p_x = list(map(lambda n: n / total_occurrences, x_counter.values()))
    s_x = entropy(p_x)
    if s_x == 0:
        return 1
    else:
        return (s_x - s_xy) / s_x

def conditional_entropy(x, y):
    r"""
    计算条件熵。

    Parameters
    ----------
    x : array-like
        类别型变量。
    y : array-like
        类别型变量。

    Returns
    -------
    float
        条件熵值。
    """
    y_counter = Counter(y)
    xy_counter = Counter(list(zip(x, y)))
    total_occurrences = sum(y_counter.values())
    entropy_xy = 0.0
    for xy in xy_counter.keys():
        p_xy = xy_counter[xy] / total_occurrences
        p_y = y_counter[xy[1]] / total_occurrences
        entropy_xy += p_xy * np.log(p_y / p_xy)
    return entropy_xy

def correlation_ratio(categories, measurements):
    """
    计算相关比率。

    Parameters
    ----------
    categories : array-like
        类别型变量。
    measurements : array-like
        连续型变量。

    Returns
    -------
    float
        相关比率值。
    """
    fcat, _ = pd.factorize(categories)
    cat_num = np.max(fcat) + 1
    y_avg_array = np.zeros(cat_num)
    n_array = np.zeros(cat_num)
    for i in range(0, cat_num):
        cat_measures = measurements[np.argwhere(fcat == i).flatten()]
        n_array[i] = len(cat_measures)
        y_avg_array[i] = np.average(cat_measures)
    y_total_avg = np.sum(np.multiply(y_avg_array, n_array)) / np.sum(n_array)
    numerator = np.sum(
        np.multiply(n_array, np.power(np.subtract(y_avg_array, y_total_avg), 2))
    )
    denominator = np.sum(np.power(np.subtract(measurements, y_total_avg), 2))
    if denominator == 0:
        return 0.0
    else:
        return np.sqrt(numerator / denominator)

def mixed_correlation(df, type_list=None, default_method="spearman"):
    """
    计算混合相关性矩阵。

    Parameters
    ----------
    df : pd.DataFrame
        包含不同数据类型变量的数据框。
    type_list : list, optional
        每列数据的类型列表，默认值为None。如果为None，将自动推断数据类型。
    default_method : str, optional
        连续变量之间相关性的默认计算方法，默认值为"spearman"。

    Returns
    -------
    pd.DataFrame
        相关性矩阵，包含类别变量和连续变量之间的相关性。
    """
    if type_list is None:
        type_list = [df[col].dtype for col in df.columns]
    cols = df.columns
    mat = pd.DataFrame(index=cols, columns=cols)

    for i in range(len(cols)):
        for j in range(i, len(cols)):
            if i == j:
                mat.iloc[i, j] = 1.0
            else:
                if type_list[i] == "object" and type_list[j] == "object":
                    mat.iloc[i, j] = cramers_v(
                        df[cols[i]], df[cols[j]])
                elif type_list[i] == "object" and type_list[j] != "object":
                    mat.iloc[i, j] = correlation_ratio(
                        df[cols[i]], df[cols[j]])
                elif type_list[i] != "object" and type_list[j] == "object":
                    mat.iloc[i, j] = correlation_ratio(
                        df[cols[j]], df[cols[i]])
                else:
                    mat.iloc[i, j] = df[cols[i]].corr(
                        df[cols[j]], method=default_method)
                mat.iloc[j, i] = mat.iloc[i, j]
    return mat.astype("float64")